{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2787d5f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import stanza\n",
    "from stanza.server.semgrex import Semgrex\n",
    "from stanza.models.common.constant import is_right_to_left\n",
    "import spacy\n",
    "from spacy import displacy\n",
    "from spacy.tokens import Doc\n",
    "from IPython.display import display, HTML\n",
    "\n",
    "\n",
    "\"\"\"\n",
    "IMPORTANT: For the code in this module to run, you must have corenlp and Java installed on your machine. Additionally,\n",
    "set an environment variable CLASSPATH equal to the path of your corenlp directory.\n",
    "\n",
    "Example: CLASSPATH=C:\\\\Users\\\\Alex\\\\PycharmProjects\\\\pythonProject\\\\stanford-corenlp-4.5.0\\\\stanford-corenlp-4.5.0\\\\*\n",
    "\"\"\"\n",
    "\n",
    "%env CLASSPATH=C:\\\\stanford-corenlp-4.5.2\\\\stanford-corenlp-4.5.2\\\\*\n",
    "def get_sentences_html(doc, language):\n",
    "    \"\"\"\n",
    "    Returns a list of the HTML strings of the dependency visualizations of a given stanza doc object.\n",
    "\n",
    "    The 'language' arg is the two-letter language code for the document to be processed.\n",
    "\n",
    "    First converts the stanza doc object to a spacy doc object and uses displacy to generate an HTML\n",
    "    string for each sentence of the doc object.\n",
    "    \"\"\"\n",
    "    html_strings = []\n",
    "\n",
    "    # blank model - we don't use any of the model features, just the visualization\n",
    "    nlp = spacy.blank(\"en\")\n",
    "    sentences_to_visualize = []\n",
    "    for sentence in doc.sentences:\n",
    "        words, lemmas, heads, deps, tags = [], [], [], [], []\n",
    "        if is_right_to_left(language):  # order of words displayed is reversed, dependency arcs remain intact\n",
    "            sent_len = len(sentence.words)\n",
    "            for word in reversed(sentence.words):\n",
    "                words.append(word.text)\n",
    "                lemmas.append(word.lemma)\n",
    "                deps.append(word.deprel)\n",
    "                tags.append(word.upos)\n",
    "                if word.head == 0:  # spaCy head indexes are formatted differently than that of Stanza\n",
    "                    heads.append(sent_len - word.id)\n",
    "                else:\n",
    "                    heads.append(sent_len - word.head)\n",
    "        else:  # left to right rendering\n",
    "            for word in sentence.words:\n",
    "                words.append(word.text)\n",
    "                lemmas.append(word.lemma)\n",
    "                deps.append(word.deprel)\n",
    "                tags.append(word.upos)\n",
    "                if word.head == 0:\n",
    "                    heads.append(word.id - 1)\n",
    "                else:\n",
    "                    heads.append(word.head - 1)\n",
    "        document_result = Doc(nlp.vocab, words=words, lemmas=lemmas, heads=heads, deps=deps, pos=tags)\n",
    "        sentences_to_visualize.append(document_result)\n",
    "\n",
    "    for line in sentences_to_visualize:  # render all sentences through displaCy\n",
    "        html_strings.append(displacy.render(line, style=\"dep\",\n",
    "                                            options={\"compact\": True, \"word_spacing\": 30, \"distance\": 100,\n",
    "                                                     \"arrow_spacing\": 20}, jupyter=False))\n",
    "    return html_strings\n",
    "\n",
    "\n",
    "def find_nth(haystack, needle, n):\n",
    "    \"\"\"\n",
    "    Returns the starting index of the nth occurrence of the substring 'needle' in the string 'haystack'.\n",
    "    \"\"\"\n",
    "    start = haystack.find(needle)\n",
    "    while start >= 0 and n > 1:\n",
    "        start = haystack.find(needle, start + len(needle))\n",
    "        n -= 1\n",
    "    return start\n",
    "\n",
    "\n",
    "def round_base(num, base=10):\n",
    "    \"\"\"\n",
    "    Rounding a number to its nearest multiple of the base. round_base(49.2, base=50) = 50.\n",
    "    \"\"\"\n",
    "    return base * round(num/base)\n",
    "\n",
    "\n",
    "def process_sentence_html(orig_html, semgrex_sentence):\n",
    "    \"\"\"\n",
    "    Takes a semgrex sentence object and modifies the HTML of the original sentence's deprel visualization,\n",
    "    highlighting words involved in the search queries and adding the label of the word inside of the semgrex match.\n",
    "\n",
    "    Returns the modified html string of the sentence's deprel visualization.\n",
    "    \"\"\"\n",
    "    tracker = {}  # keep track of which words have multiple labels\n",
    "    DEFAULT_TSPAN_COUNT = 2  # the original displacy html assigns two <tspan> objects per <text> object\n",
    "    CLOSING_TSPAN_LEN = 8  # </tspan> is 8 chars long\n",
    "    colors = ['#4477AA', '#66CCEE', '#228833', '#CCBB44', '#EE6677', '#AA3377', '#BBBBBB']\n",
    "    css_bolded_class = \"<style> .bolded{font-weight: bold;} </style>\\n\"\n",
    "    found_index = orig_html.find(\"\\n\")  # returns index where the opening <svg> ends\n",
    "    # insert the new style class into html string\n",
    "    orig_html = orig_html[: found_index + 1] + css_bolded_class + orig_html[found_index + 1:]\n",
    "\n",
    "    # Add color to words in the match, bold words in the match\n",
    "    for query in semgrex_sentence.result:\n",
    "        for i, match in enumerate(query.match):\n",
    "            color = colors[i]\n",
    "            paired_dy = 2\n",
    "            for node in match.node:\n",
    "                name, match_index = node.name, node.matchIndex\n",
    "                # edit existing <tspan> to change color and bold the text\n",
    "                start = find_nth(orig_html, \"<text\", match_index)  # finds start of svg <text> of interest\n",
    "                if match_index not in tracker:  # if we've already bolded and colored, keep the first color\n",
    "                    tspan_start = orig_html.find(\"<tspan\",\n",
    "                                                 start)  # finds start of the first svg <tspan> inside of the <text>\n",
    "                    tspan_end = orig_html.find(\"</tspan>\", start)  # finds start of the end of the above <tspan>\n",
    "                    tspan_substr = orig_html[tspan_start: tspan_end + CLOSING_TSPAN_LEN + 1] + \"\\n\"\n",
    "                    # color words in the hit and bold words in the hit\n",
    "                    edited_tspan = tspan_substr.replace('class=\"displacy-word\"', 'class=\"bolded\"').replace(\n",
    "                        'fill=\"currentColor\"', f'fill=\"{color}\"')\n",
    "                    # insert edited <tspan> object into html string\n",
    "                    orig_html = orig_html[: tspan_start] + edited_tspan + orig_html[tspan_end + CLOSING_TSPAN_LEN + 2:]\n",
    "                    tracker[match_index] = DEFAULT_TSPAN_COUNT\n",
    "\n",
    "                # next, we have to insert the new <tspan> object for the label\n",
    "                # Copy old <tspan> to copy formatting when creating new <tspan> later\n",
    "                prev_tspan_start = find_nth(orig_html[start:], \"<tspan\",\n",
    "                                            tracker[match_index] - 1) + start  # find the previous <tspan> start index\n",
    "                prev_tspan_end = find_nth(orig_html[start:], \"</tspan>\",\n",
    "                                          tracker[match_index] - 1) + start  # find the prev </tspan> start index\n",
    "                prev_tspan = orig_html[prev_tspan_start: prev_tspan_end + CLOSING_TSPAN_LEN + 1]\n",
    "\n",
    "                # Find spot to insert new tspan\n",
    "                closing_tspan_start = find_nth(orig_html[start:], \"</tspan>\", tracker[match_index]) + start\n",
    "                up_to_new_tspan = orig_html[: closing_tspan_start + CLOSING_TSPAN_LEN + 1]\n",
    "                rest_need_add_newline = orig_html[closing_tspan_start + CLOSING_TSPAN_LEN + 1:]\n",
    "\n",
    "                # Calculate proper x value in svg\n",
    "                x_value_start = prev_tspan.find('x=\"')\n",
    "                x_value_end = prev_tspan[x_value_start + 3:].find('\"') + 3  # 3 is the length of the 'x=\"' substring\n",
    "                x_value = prev_tspan[x_value_start + 3: x_value_end + x_value_start]\n",
    "\n",
    "                # Calculate proper y value in svg\n",
    "                DEFAULT_DY_VAL, dy = 2, 2\n",
    "                if paired_dy != DEFAULT_DY_VAL and node == match.node[\n",
    "                    1]:  # we're on the second node and need to adjust height to match the paired node\n",
    "                    dy = paired_dy\n",
    "                if node == match.node[0]:\n",
    "                    paired_node_level = 2\n",
    "                    if match.node[1].matchIndex in tracker:  # check if we need to adjust heights of labels\n",
    "                        paired_node_level = tracker[match.node[1].matchIndex]\n",
    "                        dif = tracker[match_index] - paired_node_level\n",
    "                        if dif > 0:  # current node has more labels\n",
    "                            paired_dy = DEFAULT_DY_VAL * dif + 1\n",
    "                            dy = DEFAULT_DY_VAL\n",
    "                        else:  # paired node has more labels, adjust this label down\n",
    "                            dy = DEFAULT_DY_VAL * (abs(dif) + 1)\n",
    "                            paired_dy = DEFAULT_DY_VAL\n",
    "\n",
    "                # Insert new <tspan> object\n",
    "                new_tspan = f'  <tspan class=\"displacy-word\" dy=\"{dy}em\" fill=\"{color}\" x={x_value}>{name[: 3].title()}.</tspan>\\n'  # abbreviate label names to 3 chars\n",
    "                orig_html = up_to_new_tspan + new_tspan + rest_need_add_newline\n",
    "                tracker[match_index] += 1\n",
    "    return orig_html\n",
    "\n",
    "\n",
    "def render_html_strings(edited_html_strings):\n",
    "    \"\"\"\n",
    "    Renders the HTML to make the edits visible\n",
    "    \"\"\"\n",
    "    for html_string in edited_html_strings:\n",
    "        display(HTML(html_string))\n",
    "\n",
    "\n",
    "def visualize_search_doc(doc, semgrex_queries, lang_code, start_match=0, end_match=10):\n",
    "    \"\"\"\n",
    "    Visualizes the semgrex results of running semgrex search on a stanza doc object with the given list of\n",
    "    semgrex queries. Returns a list of the edited HTML strings from the doc. Each element in the list represents\n",
    "    the HTML to render one of the sentences in the document.\n",
    "\n",
    "    'lang_code' is the two-letter language abbreviation for the language that the stanza doc object is written in.\n",
    "\n",
    "\n",
    "    'start_match' and 'end_match' determine which matches to visualize. Works similar to splices, so that\n",
    "    start_match=0 and end_match=10 will display the first 10 semgrex matches.\n",
    "    \"\"\"\n",
    "    matches_count = 0  # Limits number of visualizations\n",
    "    with Semgrex(classpath=\"$CLASSPATH\") as sem:\n",
    "        edited_html_strings = []\n",
    "        semgrex_results = sem.process(doc, *semgrex_queries)\n",
    "        # one html string for each sentence\n",
    "        unedited_html_strings = get_sentences_html(doc, lang_code)\n",
    "        for i in range(len(unedited_html_strings)):\n",
    "\n",
    "            if matches_count >= end_match:  # we've collected enough matches, stop early\n",
    "                break\n",
    "\n",
    "            # check if sentence has matches, if not then do not visualize\n",
    "            has_none = True\n",
    "            for query in semgrex_results.result[i].result:\n",
    "                for match in query.match:\n",
    "                    if match:\n",
    "                        has_none = False\n",
    "\n",
    "            # Process HTML if queries have matches\n",
    "            if not has_none:\n",
    "                if start_match <= matches_count < end_match:\n",
    "                    edited_string = process_sentence_html(unedited_html_strings[i], semgrex_results.result[i])\n",
    "                    edited_string = adjust_dep_arrows(edited_string)\n",
    "                    edited_html_strings.append(edited_string)\n",
    "                matches_count += 1\n",
    "\n",
    "        render_html_strings(edited_html_strings)\n",
    "    return edited_html_strings\n",
    "\n",
    "\n",
    "def visualize_search_str(text, semgrex_queries, lang_code):\n",
    "    \"\"\"\n",
    "    Visualizes the deprel of the semgrex results from running semgrex search on a string with the given list of\n",
    "    semgrex queries. Returns a list of the edited HTML strings. Each element in the list represents\n",
    "    the HTML to render one of the sentences in the document.\n",
    "\n",
    "    Internally, this function converts the string into a stanza doc object before processing the doc object.\n",
    "\n",
    "    'lang_code' is the two-letter language abbreviation for the language that the stanza doc object is written in.\n",
    "    \"\"\"\n",
    "    nlp = stanza.Pipeline(lang_code, processors=\"tokenize, pos, lemma, depparse\")\n",
    "    doc = nlp(text)\n",
    "    return visualize_search_doc(doc, semgrex_queries, lang_code)\n",
    "\n",
    "\n",
    "def adjust_dep_arrows(raw_html):\n",
    "    \"\"\"\n",
    "    The default spaCy dependency visualization has misaligned arrows.\n",
    "    We fix arrows by aligning arrow ends and bodies to the word that they are directed to. If a word has an\n",
    "    arrowhead that is pointing not directly on the word's center, align the arrowhead to match the center of the word.\n",
    "\n",
    "    returns the edited html with fixed arrow placement\n",
    "    \"\"\"\n",
    "    HTML_ARROW_BEGINNING = '<g class=\"displacy-arrow\">'\n",
    "    HTML_ARROW_ENDING = \"</g>\"\n",
    "    HTML_ARROW_ENDING_LEN = 6   # there are 2 newline chars after the arrow ending\n",
    "    arrows_start_idx = find_nth(haystack=raw_html, needle='<g class=\"displacy-arrow\">', n=1)\n",
    "    words_html, arrows_html = raw_html[: arrows_start_idx], raw_html[arrows_start_idx:]  # separate html for words and arrows\n",
    "    final_html = words_html  # continually concatenate to this after processing each arrow\n",
    "    arrow_number = 1  # which arrow we're editing (1-indexed)\n",
    "    start_idx, end_of_class_idx = find_nth(haystack=arrows_html, needle=HTML_ARROW_BEGINNING, n=arrow_number), find_nth(arrows_html, HTML_ARROW_ENDING, arrow_number)\n",
    "    while start_idx != -1:  # edit every arrow\n",
    "        arrow_section = arrows_html[start_idx: end_of_class_idx + HTML_ARROW_ENDING_LEN]  # slice a single svg arrow object\n",
    "        if arrow_section[-1] == \"<\":   # this is the last arrow in the HTML, don't cut the splice early\n",
    "            arrow_section = arrows_html[start_idx:]\n",
    "        edited_arrow_section = edit_dep_arrow(arrow_section)\n",
    "\n",
    "        final_html = final_html + edited_arrow_section  # continually update html with new arrow html until done\n",
    "\n",
    "        # Prepare for next iteration\n",
    "        arrow_number += 1\n",
    "        start_idx = find_nth(arrows_html, '<g class=\"displacy-arrow\">', n=arrow_number)\n",
    "        end_of_class_idx = find_nth(arrows_html, \"</g>\", arrow_number)\n",
    "    return final_html\n",
    "\n",
    "\n",
    "def edit_dep_arrow(arrow_html):\n",
    "    \"\"\"\n",
    "    The formatting of a displacy arrow in svg is the following:\n",
    "    <g class=\"displacy-arrow\">\n",
    "        <path class=\"displacy-arc\" id=\"arrow-c628889ffbf343e3848193a08606f10a-0-0\" stroke-width=\"2px\" d=\"M70,352.0 C70,177.0 390.0,177.0 390.0,352.0\" fill=\"none\" stroke=\"currentColor\"/>\n",
    "        <text dy=\"1.25em\" style=\"font-size: 0.8em; letter-spacing: 1px\">\n",
    "            <textPath xlink:href=\"#arrow-c628889ffbf343e3848193a08606f10a-0-0\" class=\"displacy-label\" startOffset=\"50%\" side=\"left\" fill=\"currentColor\" text-anchor=\"middle\">csubj</textPath>\n",
    "        </text>\n",
    "        <path class=\"displacy-arrowhead\" d=\"M70,354.0 L62,342.0 78,342.0\" fill=\"currentColor\"/>\n",
    "    </g>\n",
    "\n",
    "    We edit the 'd = ...' parts of the <path class ...> section to fix the arrow direction and length\n",
    "\n",
    "    returns the arrow_html with distances fixed\n",
    "    \"\"\"\n",
    "    WORD_SPACING = 50   # words start at x=50 and are separated by 100s so their x values are multiples of 50\n",
    "    M_OFFSET = 4  # length of 'd=\"M' that we search for to extract the number from d=\"M70, for instance\n",
    "    ARROW_PIXEL_SIZE = 4\n",
    "    first_d_idx, second_d_idx = find_nth(arrow_html, 'd=\"M', 1), find_nth(arrow_html, 'd=\"M', 2)  # find where d=\"M starts\n",
    "    first_d_cutoff, second_d_cutoff = arrow_html.find(\",\", first_d_idx), arrow_html.find(\",\", second_d_idx)  # isolate the number after 'M' e.g. 'M70'\n",
    "    # gives svg x values of arrow body starting position and arrowhead position\n",
    "    arrow_position, arrowhead_position = float(arrow_html[first_d_idx + M_OFFSET: first_d_cutoff]), float(arrow_html[second_d_idx + M_OFFSET: second_d_cutoff])\n",
    "    # gives starting index of where 'fill=\"none\"' or 'fill=\"currentColor\"' begin, reference points to end the d= section\n",
    "    first_fill_start_idx, second_fill_start_idx = find_nth(arrow_html, \"fill\", n=1), find_nth(arrow_html, \"fill\", n=3)\n",
    "\n",
    "    # isolate the d= ... section to edit\n",
    "    first_d, second_d = arrow_html[first_d_idx: first_fill_start_idx], arrow_html[second_d_idx: second_fill_start_idx]\n",
    "    first_d_split, second_d_split = first_d.split(\",\"), second_d.split(\",\")\n",
    "\n",
    "    if arrow_position == arrowhead_position:  # This arrow is incoming onto the word, center the arrow/head to word center\n",
    "        corrected_arrow_pos = corrected_arrowhead_pos = round_base(arrow_position, base=WORD_SPACING)\n",
    "\n",
    "        # edit first_d  -- arrow body\n",
    "        second_term = first_d_split[1].split(\" \")[0] + \" \" + str(corrected_arrow_pos)\n",
    "        first_d = 'd=\"M' + str(corrected_arrow_pos) + \",\" + second_term + \",\" + \",\".join(first_d_split[2:])\n",
    "\n",
    "        # edit second_d  -- arrowhead\n",
    "        second_term = second_d_split[1].split(\" \")[0] + \" L\" + str(corrected_arrowhead_pos - ARROW_PIXEL_SIZE)\n",
    "        third_term = second_d_split[2].split(\" \")[0] + \" \" + str(corrected_arrowhead_pos + ARROW_PIXEL_SIZE)\n",
    "        second_d = 'd=\"M' + str(corrected_arrowhead_pos) + \",\" + second_term + \",\" + third_term + \",\" + \",\".join(second_d_split[3:])\n",
    "    else:  # This arrow is outgoing to another word, center the arrow/head to that word's center\n",
    "        corrected_arrowhead_pos = round_base(arrowhead_position, base=WORD_SPACING)\n",
    "\n",
    "        # edit first_d -- arrow body\n",
    "        third_term = first_d_split[2].split(\" \")[0] + \" \" + str(corrected_arrowhead_pos)\n",
    "        fourth_term = first_d_split[3].split(\" \")[0] + \" \" + str(corrected_arrowhead_pos)\n",
    "        terms = [first_d_split[0], first_d_split[1], third_term, fourth_term] + first_d_split[4:]\n",
    "        first_d = \",\".join(terms)\n",
    "\n",
    "        # edit second_d -- arrow head\n",
    "        first_term = f'd=\"M{corrected_arrowhead_pos}'\n",
    "        second_term = second_d_split[1].split(\" \")[0] + \" L\" + str(corrected_arrowhead_pos - ARROW_PIXEL_SIZE)\n",
    "        third_term = second_d_split[2].split(\" \")[0] + \" \" + str(corrected_arrowhead_pos + ARROW_PIXEL_SIZE)\n",
    "        terms = [first_term, second_term, third_term] + second_d_split[3:]\n",
    "        second_d = \",\".join(terms)\n",
    "    # rebuild and return html\n",
    "    return arrow_html[:first_d_idx] + first_d + \" \" + arrow_html[first_fill_start_idx:second_d_idx] + second_d + \" \" + arrow_html[second_fill_start_idx:]\n",
    "\n",
    "\n",
    "def main():\n",
    "    nlp = stanza.Pipeline(\"en\", processors=\"tokenize,pos,lemma,depparse\")\n",
    "\n",
    "    # doc = nlp(\"This a dummy sentence. Banning opal removed all artifact decks from the meta.  I miss playing lantern. This is a dummy sentence.\")\n",
    "    doc = nlp(\"Banning opal removed artifact decks from the meta. Banning tennis resulted in players banning people.\")\n",
    "    # A single result .result[i].result[j] is a list of matches for sentence i on semgrex query j.\n",
    "    queries = [\"{pos:NN}=object <obl {}=action\",\n",
    "               \"{cpos:NOUN}=thing <obj {cpos:VERB}=action\"]\n",
    "    res = visualize_search_doc(doc, queries, \"en\")\n",
    "    print(res[0])  # see the first sentence's deprel visualization HTML\n",
    "    print(\"---------------------------------------\")\n",
    "    print(res[1])  # second sentence's deprel visualization HTML\n",
    "    return\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    main()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
